package com.xcj.springai.chatMemory;

import com.baomidou.mybatisplus.extension.conditions.query.LambdaQueryChainWrapper;
import com.xcj.springai.entity.ChatConversationEntity;
import com.xcj.springai.entity.ChatMessageEntity;
import com.xcj.springai.service.ChatConversationService;
import com.xcj.springai.service.ChatMessageService;
import com.xcj.springai.utils.MessageConverter;
import jakarta.annotation.Resource;
import lombok.Builder;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

import java.util.Collections;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;


/**
 *将对话消息 持久化存入 MySQL 数据库，
 * 并在需要时读取出来，提供给大模型形成上下文记忆。
 */
@Component
@Builder
public class MySqlChatMemory implements ChatMemory {

    private static final int DEFAULT_MAX_MESSAGES = 20;
    @Resource
    private ChatMessageService chatMessageService;

    @Resource
    private ChatConversationService chatConversationService;

    /**
     * 保存新对话消息到数据库中
     *
     * @param conversationId
     * @param messages
     */
    @Override
    @Transactional(rollbackFor = Exception.class)
    public void add(String conversationId, List<Message> messages) {
        // 分割 conversationId
        String[] split = conversationId.split(":");
        // 需要在这里获取到用户的userId，然后在保存消息的时候携带对应的userId
        // 将消息列表转换为 ChatMessage 列表
        List<ChatMessageEntity> chatMessageList = messages.stream()
//                .map(message -> MessageConverter.toChatMessage(message, split[0],split[1], split[2]))
                .map(message -> MessageConverter.toChatMessage(message, split[0],split[1]))
                .collect(Collectors.toList());

        chatMessageList.forEach(v -> {
            MessageType roleType = v.getRoleType();
            if (roleType.equals(MessageType.ASSISTANT)) {
                if (split[0].equals("TOOL")) {
                    v.setMt("tool_call");
                } else {
                    v.setMt(null);
                }
            }
        });

        // 判断第一条消息的类型是否为 USER
        if (messages.get(0).getMessageType() == MessageType.USER) {
            // 创建 ChatConversation 对象
            ChatConversationEntity chatConversation = new ChatConversationEntity();
            // 设置 ChatConversation 的 ID
            chatConversation.setId(split[1]);
            // 设置 ChatConversation 的 UserId
//            chatConversation.setUserId(split[2]);
            // 设置 ChatConversation 的标题为第一条消息的文本内容
            chatConversation.setTitle(messages.get(0).getText());
            // 设置 ChatConversation 的创建时间
            chatConversation.setCreateTime(new Date());
            // 保存或更新 ChatConversation
            chatConversationService.saveOrUpdate(chatConversation);
        }
        // 批量保存 ChatMessage 列表
        chatMessageService.saveBatch(chatMessageList, chatMessageList.size());
    }


    /**
     * 未拆分的chat
     * @param conversationId
     * @return
     */

    @Override
    public List<Message> get(String conversationId) {
        String[] split = conversationId.split(":");
        List<ChatMessageEntity> list = chatMessageService.lambdaQuery()
                .eq(ChatMessageEntity::getConversationId, split[0])
                .eq(ChatMessageEntity::getUserId, split[1])
                .isNull(ChatMessageEntity::getMt)
//                .eq(ChatMessageEntity::getRoleType, "USER")
//                .orderByDesc(ChatMessageEntity::getCreateTime)
                .last(DEFAULT_MAX_MESSAGES > 0, "limit " + DEFAULT_MAX_MESSAGES)
                .list();

        return list.stream()
                .map(MessageConverter::toMessage)
                .collect(Collectors.toList());
    }

//    @Override
//    public List<Message> get(String conversationId) {
//        String[] parts = conversationId.split(":");
//        if (parts.length != 3) {
//            throw new IllegalArgumentException("Invalid conversationId format: " + conversationId);
//        }
//        String prefix = parts[0];
//        String convId = parts[1];
//        String userId = parts[2];
//
//        boolean isRagConversation = "RAG".equalsIgnoreCase(prefix);
//        LambdaQueryChainWrapper<ChatMessageEntity> query = chatMessageService.lambdaQuery();
//        query.eq(ChatMessageEntity::getConversationId, convId)
//                .eq(ChatMessageEntity::getUserId, userId);
//
//        //根据前缀来筛选数据，如果是RAG对话则mt字段为空，否则mt字段不为空
//        if (isRagConversation) {
//            query.isNull(ChatMessageEntity::getMt);
//        } else {
//            query.isNotNull(ChatMessageEntity::getMt);
//        }
//
//        if (DEFAULT_MAX_MESSAGES > 0) {
//            query.last("limit " + DEFAULT_MAX_MESSAGES);
//        }
//
//        // 获取数据列表
//        List<ChatMessageEntity> list = query.list();
//
//        // 转换为业务消息对象
//        return list.stream()
//                .map(MessageConverter::toMessage)
//                .collect(Collectors.toList());
//    }



    @Override
    public void clear(String conversationId) {

    }
}